from pydantic import SecretStr
from sentence_transformers import SentenceTransformer
from typing import List, Dict, Optional
import logging
import uuid
# services/vector_db.py
from chromadb.config import Settings
from chromadb import HttpClient

from config import settings

logger = logging.getLogger(__name__)


class VectorDBManager:
    def __init__(self):
        # 解析认证信息
        auth_parts = settings.chroma_auth.get_secret_value().split(":")
        if len(auth_parts) != 2:
            raise ValueError("CHROMA_AUTH配置格式不正确，应为'username:password'")

        username, password = auth_parts[0], auth_parts[1]
        print(username,password)
        # 配置SSL上下文（支持自签名证书）
        # ssl_context = ssl.create_default_context()
        # if not settings.chroma_ssl_verify:
        #     ssl_context.check_hostname = False
        #     ssl_context.verify_mode = ssl.CERT_NONE

        # 初始化客户端
        self.client = HttpClient(
            host=settings.chroma_host,
            port=settings.chroma_port
        )
        self.collection = self.client.get_or_create_collection("knowledge_base")
        self.embedder = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")

    def upsert_chunks(self, chunks: List[Dict]):
        """插入或更新文档块"""
        try:
            ids = [str(uuid.uuid4()) for _ in chunks]
            embeddings = self._generate_embeddings(
                [c["content"] for c in chunks]
            )

            self.collection.upsert(
                ids=ids,
                embeddings=embeddings,
                documents=[c["content"] for c in chunks],
                metadatas=[c["metadata"] for c in chunks]
            )
            logger.info(f"成功插入 {len(chunks)} 个文档块")
        except Exception as e:
            logger.error(f"插入文档块失败: {str(e)}")

    def search(self, query: str, top_k: int = 5) -> List[Dict]:
        """向量搜索"""
        try:
            query_embedding = self._generate_embeddings([query])
            results = self.collection.query(
                query_embeddings=query_embedding,
                n_results=top_k
            )
            return self._format_results(results)
        except Exception as e:
            logger.error(f"搜索失败: {str(e)}")
            return []

    def update_chunk(self, chunk_id: str, new_content: str):
        """更新文档块"""
        try:
            self.collection.update(
                ids=chunk_id,
                documents=new_content,
                embeddings=self._generate_embeddings([new_content])
            )
        except Exception as e:
            logger.error(f"更新块 {chunk_id} 失败: {str(e)}")

    def _generate_embeddings(self, texts: List[str]) -> List[List[float]]:
        return self.embedder.encode(texts).tolist()

    def _format_results(self, results) -> List[Dict]:
        return [{
            "id": results["ids"][0][i],
            "content": results["documents"][0][i],
            "metadata": results["metadatas"][0][i],
            "score": results["distances"][0][i]
        } for i in range(len(results["ids"][0]))]

